#include "stdafx.h"
#include "RBTree.h"

namespace WTB{

	CRBTreeNode* CRBTreeNode::GetNextNode()
	{
		if (GetParent() == this)
		{
			return NULL;
		}

		CRBTreeNode *pParent;
		CRBTreeNode* pNode = this;

		/* If we have a right-hand child, go down and then left as far
		as we can. */

		if (pNode->m_pRightChild)
		{
			pNode = pNode->m_pRightChild;
			while(pNode->m_pLeftChild)
				pNode = pNode->m_pLeftChild;
			return pNode;
		}

		/* No right-hand children.  Everything down and left is
		smaller than us, so any 'next' node must be in the general
		direction of our parent. Go up the tree; any time the
		ancestor is a right-hand child of its parent, keep going
		up. First time it's a left-hand child of its parent, said
		parent is our 'next' node. */
		while ((pParent = pNode->GetParent()) && pNode == pParent->m_pRightChild)
			pNode = pParent;

		return pParent;
	}//GetNextNode
	CRBTreeNode* CRBTreeNode::GetPreNode()
	{
		if (GetParent() == this)
		{
			return NULL;
		}

		CRBTreeNode *pParent;
		CRBTreeNode* pNode = this;

		/* If we have a left-hand child, go down and then right as far
		as we can. */
		if (pNode->m_pLeftChild) {
			pNode = pNode->m_pLeftChild; 
			while (pNode->m_pRightChild)
				pNode=pNode->m_pRightChild;
			return pNode;
		}

		/* No left-hand children. Go up till we find an ancestor which
		is a right-hand child of its parent */
		while ((pParent = pNode->GetParent()) && pNode == pParent->m_pLeftChild)
			pNode = pParent;

		return pParent;
	}

	void CRBTree::_RotateLeft(CRBTreeNode* _pRotateNode)
	{
		CRBTreeNode *pRight = _pRotateNode->m_pRightChild;
		CRBTreeNode *pParent = _pRotateNode->GetParent();

		if ((_pRotateNode->m_pRightChild = pRight->m_pLeftChild))
		{
			pRight->SetParent(_pRotateNode);
		}
		pRight->m_pLeftChild = _pRotateNode;

		pRight->SetParent(pParent);

		if (pParent)
		{
			if (_pRotateNode == pParent->m_pLeftChild)
			{
				pParent->m_pLeftChild = pRight;
			}
			else
			{
				pParent->m_pRightChild = pRight;
			}
		}
		else
		{
			m_pRBTreeRoot = pRight;
		}
		_pRotateNode->SetParent(pRight);
	}//end _RotateLeft

	void CRBTree::_RotateRight(CRBTreeNode* _pRotateNode)
	{
		CRBTreeNode *pLeft = _pRotateNode->m_pLeftChild;
		CRBTreeNode *pParent = _pRotateNode->GetParent();

		if ((_pRotateNode->m_pLeftChild = pLeft->m_pRightChild))
		{
			pLeft->m_pRightChild->SetParent(_pRotateNode);
		}
		pLeft->m_pRightChild = _pRotateNode;

		pLeft->SetParent(pParent);

		if (pParent)
		{
			if (_pRotateNode == pParent->m_pRightChild)
			{
				pParent->m_pRightChild = pLeft;
			}
			else
			{
				pParent->m_pLeftChild = pLeft;
			}
		}
		else
		{
			m_pRBTreeRoot = pLeft;
		}
		_pRotateNode->SetParent(pLeft);
	}//end _RotateRight


	void CRBTree::InsertColor(CRBTreeNode *_pInsertNode)
	{
		CRBTreeNode *pParent;
		CRBTreeNode *pGrandParent;

		while((pParent = _pInsertNode->GetParent()) && pParent->IsRed())
		{
			pGrandParent = pParent->GetParent();

			if (pParent == pGrandParent->m_pLeftChild)
			{
				CRBTreeNode *pUncle = pGrandParent->m_pRightChild;
				if (pUncle && pUncle->IsRed())
				{
					pUncle->SetBlack();
					pParent->SetBlack();
					pGrandParent->SetRed();
					_pInsertNode = pGrandParent;
					continue;
				}

				if (pParent->m_pRightChild == _pInsertNode)
				{
					CRBTreeNode *pTmp;
					_RotateLeft(pParent);
					pTmp = pParent;
					pParent = _pInsertNode;
					_pInsertNode = pTmp;
				}

				pParent->SetBlack();
				pGrandParent->SetRed();
				_RotateRight(pGrandParent);
			}
			//end if pParent == pGrandParent->m_pLeftChild
			else
			{
				CRBTreeNode *pUncle = pGrandParent->m_pLeftChild;
				if (pUncle && pUncle->IsRed())
				{
					pUncle->SetBlack();
					pParent->SetBlack();
					pGrandParent->SetRed();
					_pInsertNode = pGrandParent;
					continue;
				}

				if (pParent->m_pLeftChild == _pInsertNode)
				{
					CRBTreeNode *pTmp;
					_RotateRight(pParent);
					pTmp = pParent;
					pParent = _pInsertNode;
					_pInsertNode = pTmp;
				}

				pParent->SetBlack();
				pGrandParent->SetRed();
				_RotateLeft(pGrandParent);
			}
			//end else
		}
		//end while

		m_pRBTreeRoot->SetBlack();
	}//end InsertColor

	void CRBTree::_EraseColor(CRBTreeNode* _pEraseNode,CRBTreeNode* _pParent)
	{
		CRBTreeNode *pOther;

		while((!_pEraseNode || _pEraseNode->IsBlack()) && _pEraseNode != m_pRBTreeRoot)
		{
			if (_pParent->m_pLeftChild == _pEraseNode)
			{
				pOther = _pParent->m_pRightChild;
				if (pOther->IsRed())
				{
					pOther->SetBlack();
					_pParent->SetRed();
					_RotateLeft(_pParent);
					pOther = _pParent->m_pRightChild;
				}
				if ((!pOther->m_pLeftChild || pOther->m_pLeftChild->IsBlack()) &&
					(!pOther->m_pRightChild || pOther->m_pRightChild->IsBlack()))
				{
					pOther->SetRed();
					_pEraseNode = _pParent;
					_pParent = _pEraseNode->GetParent();
				}
				else
				{
					if (!pOther->m_pRightChild || pOther->m_pRightChild->IsBlack())
					{
						pOther->m_pLeftChild->SetBlack();
						pOther->SetRed();
						_RotateRight(pOther);
						pOther = _pParent->m_pRightChild;
					}
					pOther->SetColor(_pParent->GetColor());
					_pParent->SetBlack();
					pOther->m_pRightChild->SetBlack();
					_RotateLeft(_pParent);
					_pEraseNode = m_pRBTreeRoot;
					break;
				}
			}//end if _pParent->m_pLeftChild == _pEraseNode

			else
			{
				pOther = _pParent->m_pLeftChild;
				if (pOther->IsRed())
				{
					pOther->SetBlack();
					_pParent->SetRed();
					_RotateRight(_pParent);
					pOther = _pParent->m_pLeftChild;
				}
				if ((!pOther->m_pLeftChild || pOther->m_pLeftChild->IsBlack()) &&
					(!pOther->m_pRightChild || pOther->m_pRightChild->IsBlack()))
				{
					pOther->SetRed();
					_pEraseNode = _pParent;
					_pParent = _pEraseNode->GetParent();
				}
				else
				{
					if (!pOther->m_pLeftChild || pOther->m_pLeftChild->IsBlack())
					{
						pOther->m_pRightChild->SetBlack();
						pOther->SetRed();
						_RotateLeft(pOther);
						pOther = _pParent->m_pLeftChild;
					}
					pOther->SetColor(_pParent->GetColor());
					_pParent->SetBlack();
					pOther->m_pLeftChild->SetBlack();
					_RotateRight(_pParent);
					_pEraseNode = m_pRBTreeRoot;
					break;
				}
			}// end else
		}//end while
		if (_pEraseNode)
		{
			_pEraseNode->SetBlack();
		}
	}//end _EraseColor

	void CRBTree::Erase(CRBTreeNode* _pEraseNode)
	{
		CRBTreeNode *pChild;
		CRBTreeNode *pParent;
		unsigned long ulColor;

		if (!_pEraseNode->m_pLeftChild)
		{
			pChild = _pEraseNode->m_pRightChild;
		}
		else if (!_pEraseNode->m_pRightChild)
		{
			pChild = _pEraseNode->m_pLeftChild;
		}
		else
		{
			CRBTreeNode *pOld = _pEraseNode;
			CRBTreeNode *pLeft;

			_pEraseNode = _pEraseNode->m_pRightChild;
			while((pLeft = _pEraseNode->m_pLeftChild) != NULL)
			{
				_pEraseNode = pLeft;
			}

			if (pOld->GetParent())
			{
				if (pOld->GetParent()->m_pLeftChild == pOld)
				{
					pOld->GetParent()->m_pLeftChild = _pEraseNode;
				}
				else
				{
					pOld->GetParent()->m_pRightChild = _pEraseNode;
				}
			}
			else
			{
				m_pRBTreeRoot = _pEraseNode;
			}

			pChild = _pEraseNode->m_pRightChild;
			pParent = _pEraseNode->GetParent();
			ulColor = _pEraseNode->GetColor();

			if (pParent == pOld)
			{
				pParent = _pEraseNode;
			}
			else
			{
				if (pChild)
				{
					pChild->SetParent(pParent);
				}
				pParent->m_pLeftChild = pChild;

				_pEraseNode->m_pRightChild = pOld->m_pRightChild;
				pOld->m_pRightChild->SetParent(_pEraseNode);
			}

			_pEraseNode->SetParent(pOld->GetParent());
			_pEraseNode->SetColor(pOld->GetColor());
			_pEraseNode->m_pLeftChild = pOld->m_pLeftChild;
			pOld->m_pLeftChild->SetParent(_pEraseNode);

			goto color;
		}//end else

		pParent = _pEraseNode->GetParent();
		ulColor = _pEraseNode->GetColor();

		if (pChild)
		{
			pChild->SetParent(pParent);
		}
		if (pParent)
		{
			if (pParent->m_pLeftChild == _pEraseNode)
			{
				pParent->m_pLeftChild = pChild;
			}
			else
			{
				pParent->m_pRightChild = pChild;
			}
		}
		else
		{
			m_pRBTreeRoot = pChild;
		}

color:
		if (ulColor == CRBTreeNode::BLACK)
		{
			_EraseColor(pChild,pParent);
		}
	}//end Erase

	CRBTreeNode* CRBTree::GetFirstNode()
	{
		CRBTreeNode *pNode;

		pNode = m_pRBTreeRoot;
		if (NULL != pNode)
		{
			while(pNode->m_pLeftChild)
				pNode = pNode->m_pLeftChild;
		}

		return pNode;
	}//GetFirstNode

	CRBTreeNode* CRBTree::GetLastNode()
	{
		CRBTreeNode *pNode;

		pNode = m_pRBTreeRoot;
		if (NULL != pNode)
		{
			while(pNode->m_pRightChild)
				pNode = pNode->m_pRightChild;
		}

		return pNode;
	}//GetLastNode

}